"""
AgentPrune框架：智能体通信剪枝系统
用于降低智能体间的冗余通信，提高系统效率
支持多种模型后端（vllm、ollama等）
"""
import asyncio
import json
import hashlib
import logging
from typing import Dict, Any, List, Optional, Set, Tuple
from dataclasses import dataclass, asdict
from datetime import datetime
from enum import Enum
from abc import ABC, abstractmethod
import numpy as np
from utils.logger import Logger

logger = Logger.get_logger("agent_prune")


class MessageType(Enum):
    """消息类型枚举"""
    TASK_REQUEST = "task_request"
    TASK_RESPONSE = "task_response"
    STATUS_UPDATE = "status_update"
    HEARTBEAT = "heartbeat"
    ERROR = "error"
    OTHER = "other"


class PruningStrategy(Enum):
    """剪枝策略枚举"""
    CONTENT_BASED = "content_based"  # 基于内容相似性
    FREQUENCY_BASED = "frequency_based"  # 基于频率统计
    SEMANTIC_BASED = "semantic_based"  # 基于语义嵌入
    HYBRID = "hybrid"  # 混合策略


@dataclass
class MessageMetadata:
    """消息元数据"""
    message_id: str
    sender_id: str
    receiver_id: str
    message_type: MessageType
    timestamp: float
    priority: int = 0
    importance_score: float = 0.0
    is_prunable: bool = True


@dataclass
class PruneContext:
    """剪枝上下文"""
    conversation_history: List[Dict[str, Any]]
    agent_roles: Dict[str, str]
    recent_messages: Dict[str, List[Dict[str, Any]]]  # sender -> [messages]
    model_backend: str = "vllm"  # 支持vllm、ollama等


class BasePruner(ABC):
    """基础剪枝器抽象类"""
    
    @abstractmethod
    async def should_prune(self, message: Dict[str, Any], context: PruneContext) -> Tuple[bool, float]:
        """
        判断消息是否应该被剪枝
        返回: (是否剪枝, 剪枝置信度)
        """
        pass
    
    @abstractmethod
    def get_pruning_reason(self) -> str:
        """获取剪枝原因"""
        pass


class ContentBasedPruner(BasePruner):
    """基于内容的剪枝器"""
    
    def __init__(self, similarity_threshold: float = 0.85, window_size: int = 10):
        self.similarity_threshold = similarity_threshold
        self.window_size = window_size
        self.last_reason = ""
    
    async def should_prune(self, message: Dict[str, Any], context: PruneContext) -> Tuple[bool, float]:
        """基于内容相似度判断是否剪枝"""
        if not message.get("content") or not context.recent_messages.get(message.get("sender_id")):
            return False, 0.0
        
        recent_msgs = context.recent_messages[message.get("sender_id")][-self.window_size:]
        current_content = message.get("content", "").strip()
        
        # 计算与最近消息的相似度
        for recent_msg in reversed(recent_msgs):
            recent_content = recent_msg.get("content", "").strip()
            if not recent_content:
                continue
            
            similarity = self._calculate_content_similarity(current_content, recent_content)
            if similarity > self.similarity_threshold:
                self.last_reason = f"内容与最近消息相似度高 ({similarity:.2f} > {self.similarity_threshold})"
                return True, similarity
        
        return False, 0.0
    
    def _calculate_content_similarity(self, text1: str, text2: str) -> float:
        """计算文本内容相似度"""
        # 使用简单的重叠词比例作为相似度
        if not text1 or not text2:
            return 0.0
        
        words1 = set(text1.lower().split())
        words2 = set(text2.lower().split())
        
        if not words1 or not words2:
            return 0.0
        
        # 计算Jaccard相似度
        intersection = len(words1.intersection(words2))
        union = len(words1.union(words2))
        
        return intersection / union
    
    def get_pruning_reason(self) -> str:
        return self.last_reason


class FrequencyBasedPruner(BasePruner):
    """基于频率的剪枝器"""
    
    def __init__(self, max_frequency: int = 3, time_window: int = 60):
        self.max_frequency = max_frequency
        self.time_window = time_window  # 秒
        self.last_reason = ""
        self.message_counts: Dict[str, List[float]] = {}
    
    async def should_prune(self, message: Dict[str, Any], context: PruneContext) -> Tuple[bool, float]:
        """基于消息频率判断是否剪枝"""
        content_hash = self._get_content_hash(message.get("content", ""))
        current_time = datetime.now().timestamp()
        
        # 清理过期的计数
        if content_hash in self.message_counts:
            self.message_counts[content_hash] = [
                t for t in self.message_counts[content_hash] 
                if current_time - t < self.time_window
            ]
        else:
            self.message_counts[content_hash] = []
        
        # 检查频率
        frequency = len(self.message_counts[content_hash])
        if frequency >= self.max_frequency:
            self.last_reason = f"消息频率过高 ({frequency}次/{self.time_window}秒)"
            # 置信度基于频率超过阈值的程度
            confidence = min(1.0, (frequency - self.max_frequency) / 2 + 0.8)
            return True, confidence
        
        # 记录本次消息
        self.message_counts[content_hash].append(current_time)
        return False, 0.0
    
    def _get_content_hash(self, content: str) -> str:
        """获取内容的哈希值"""
        return hashlib.md5(content.strip().lower().encode()).hexdigest()
    
    def get_pruning_reason(self) -> str:
        return self.last_reason


class HybridPruner(BasePruner):
    """混合剪枝器"""
    
    def __init__(self):
        self.content_pruner = ContentBasedPruner()
        self.frequency_pruner = FrequencyBasedPruner()
        self.last_reason = ""
    
    async def should_prune(self, message: Dict[str, Any], context: PruneContext) -> Tuple[bool, float]:
        """混合策略判断是否剪枝"""
        # 先检查频率
        freq_prune, freq_conf = await self.frequency_pruner.should_prune(message, context)
        if freq_prune:
            self.last_reason = self.frequency_pruner.get_pruning_reason()
            return True, freq_conf
        
        # 再检查内容相似度
        content_prune, content_conf = await self.content_pruner.should_prune(message, context)
        if content_prune:
            self.last_reason = self.content_pruner.get_pruning_reason()
            return True, content_conf
        
        return False, 0.0
    
    def get_pruning_reason(self) -> str:
        return self.last_reason


class AgentPruneManager:
    """AgentPrune管理器"""
    
    def __init__(self, strategy: PruningStrategy = PruningStrategy.HYBRID):
        self.pruners: Dict[PruningStrategy, BasePruner] = {
            PruningStrategy.CONTENT_BASED: ContentBasedPruner(),
            PruningStrategy.FREQUENCY_BASED: FrequencyBasedPruner(),
            PruningStrategy.SEMANTIC_BASED: self._get_semantic_pruner(),
            PruningStrategy.HYBRID: HybridPruner()
        }
        self.current_strategy = strategy
        self.message_history: Dict[str, List[Dict[str, Any]]] = {}
        self.agent_roles: Dict[str, str] = {}
        self.pruning_stats: Dict[str, Any] = {
            "total_messages": 0,
            "pruned_messages": 0,
            "pruning_rate": 0.0,
            "pruning_by_reason": {}
        }
        self.model_backend = "vllm"  # 默认使用vllm
    
    def _get_semantic_pruner(self) -> BasePruner:
        """获取语义剪枝器（可根据模型后端选择）"""
        # 这里返回一个基于内容的剪枝器作为默认实现
        # 后续可以根据不同的模型后端实现更复杂的语义剪枝
        return ContentBasedPruner(similarity_threshold=0.75)
    
    def set_model_backend(self, backend: str) -> None:
        """设置模型后端"""
        self.model_backend = backend
        logger.info(f"AgentPrune: 设置模型后端为 {backend}")
    
    async def process_message(self, message: Dict[str, Any], sender_id: str, receiver_id: str) -> Tuple[bool, Optional[str]]:
        """
        处理消息，判断是否应该剪枝
        返回: (是否传递消息, 剪枝原因)
        """
        # 记录消息统计
        self.pruning_stats["total_messages"] += 1
        
        # 更新消息历史
        if sender_id not in self.message_history:
            self.message_history[sender_id] = []
        
        # 限制历史消息数量
        self.message_history[sender_id].append({
            "content": message.get("content", ""),
            "timestamp": datetime.now().timestamp(),
            "message_type": message.get("message_type", "OTHER")
        })
        
        if len(self.message_history[sender_id]) > 50:
            self.message_history[sender_id] = self.message_history[sender_id][-50:]
        
        # 创建剪枝上下文
        context = PruneContext(
            conversation_history=[],  # 简化实现，实际可以传入完整对话历史
            agent_roles=self.agent_roles,
            recent_messages=self.message_history,
            model_backend=self.model_backend
        )
        
        # 判断是否应该剪枝
        pruner = self.pruners.get(self.current_strategy, self.pruners[PruningStrategy.HYBRID])
        should_prune, confidence = await pruner.should_prune(
            {**message, "sender_id": sender_id, "receiver_id": receiver_id},
            context
        )
        
        if should_prune and confidence > 0.5:  # 置信度阈值
            reason = pruner.get_pruning_reason()
            self._update_pruning_stats(reason)
            logger.debug(f"AgentPrune: 消息被剪枝 - {reason} (sender={sender_id}, receiver={receiver_id})")
            return False, reason
        
        return True, None
    
    def _update_pruning_stats(self, reason: str) -> None:
        """更新剪枝统计"""
        self.pruning_stats["pruned_messages"] += 1
        self.pruning_stats["pruning_rate"] = (
            self.pruning_stats["pruned_messages"] / self.pruning_stats["total_messages"]
        )
        
        if reason not in self.pruning_stats["pruning_by_reason"]:
            self.pruning_stats["pruning_by_reason"][reason] = 0
        self.pruning_stats["pruning_by_reason"][reason] += 1
    
    def register_agent_role(self, agent_id: str, role: str) -> None:
        """注册智能体角色"""
        self.agent_roles[agent_id] = role
    
    def get_pruning_stats(self) -> Dict[str, Any]:
        """获取剪枝统计信息"""
        return self.pruning_stats.copy()
    
    def set_pruning_strategy(self, strategy: PruningStrategy) -> None:
        """设置剪枝策略"""
        if strategy in self.pruners:
            self.current_strategy = strategy
            logger.info(f"AgentPrune: 设置剪枝策略为 {strategy.value}")
        else:
            logger.warning(f"AgentPrune: 未知的剪枝策略 {strategy.value}")


# 创建全局实例
try:
    agent_prune_manager = AgentPruneManager()
except Exception as e:
    logger.error(f"AgentPruneManager初始化失败: {str(e)}")
    # 降级处理：创建一个简单的管理器
    class SimplePruneManager:
        async def process_message(self, *args, **kwargs):
            return True, None
        def register_agent_role(self, *args, **kwargs): pass
        def get_pruning_stats(self): return {}
        def set_model_backend(self, backend): pass
    
    agent_prune_manager = SimplePruneManager()


# 测试函数（供开发使用）
async def test_agent_prune():
    """测试AgentPrune功能"""
    manager = AgentPruneManager()
    
    # 测试内容重复剪枝
    message1 = {"content": "这是一条测试消息", "message_type": "task_request"}
    message2 = {"content": "这是一条测试消息", "message_type": "task_request"}
    
    # 第一条消息应该通过
    pass1, reason1 = await manager.process_message(message1, "agent1", "agent2")
    print(f"测试1: 消息通过={pass1}, 原因={reason1}")
    
    # 第二条消息应该被剪枝
    pass2, reason2 = await manager.process_message(message2, "agent1", "agent2")
    print(f"测试2: 消息通过={pass2}, 原因={reason2}")
    
    # 测试频率剪枝
    for i in range(5):
        msg = {"content": f"频繁消息{i}", "message_type": "status_update"}
        result, reason = await manager.process_message(msg, "agent3", "agent4")
        print(f"频率测试{i+3}: 通过={result}, 原因={reason}")
    
    print("剪枝统计:", manager.get_pruning_stats())


if __name__ == "__main__":
    # 运行测试
    asyncio.run(test_agent_prune())
